import os
from os.path import join
import numpy as np
import cv2
import matplotlib.pyplot as plt
from side_extractor import process_piece, plot_side_images
from functools import partial
import traceback
filenames = os.listdir('images')
filenames.sort()
def plot_grid(size, out_dict, *image_keys):
h, w = size
for idx, img_key in enumerate(image_keys, start=1):
plt.subplot(h * 100 + w * 10 + idx)
if img_key[0] == '_':
plt.imshow(out_dict[img_key[1:]], cmap='gray')
else:
plt.imshow(out_dict[img_key])
label_tuples = [('A', 74), ('B', 43), ('C', 19), ('D', 72), ('E', 11)]
def create_label(label_tuple):
letter, max_num = label_tuple
for i in range(1, max_num + 1):
label = letter + str(i) if i >= 10 else letter + '0' + str(i)
yield label
labels = []
for label_tuple in label_tuples:
for label in create_label(label_tuple):
labels.append(label)
postprocess = partial(cv2.blur, ksize=(3, 3))
results = []
for filename, label in zip(filenames, labels):
img = cv2.imread(join('images', filename))
img = img[1750:2500, 1000:2000]
out_dict = process_piece(img, after_segmentation_func=postprocess)
if 'error' in out_dict:
print(out_dict['error'])
traceback.print_exc()
continue
plt.figure(figsize=(6, 6))
plt.title("{0} - {1}".format(filename, label))
plt.imshow(out_dict['class_image'])
#plot_grid((3, 3), out_dict, '_segmented', '_extracted', '_edges', 'class_image')
plt.show()
plot_side_images(out_dict['side_images'], out_dict['inout'])
results.append({'side_images': out_dict['side_images'], 'inout': out_dict['inout']})
to_ignore = ['D11', 'D70']
for label, result in zip(labels, results):
if label in to_ignore:
continue
for i, (side_image, io) in enumerate(zip(result['side_images'], result['inout']), start=1):
out_io = 'int' if io == 'in' else 'out'
side_image = side_image * 255
out_filename = "{0}_{1}_{2}.jpg".format(label, i, out_io)
out_path = join('sides', out_filename)
cv2.imwrite(out_path, side_image)